package snailAI.Config;
import org.apache.http.HttpHost;
import org.elasticsearch.client.RestClient;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions;
import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
@Configuration
public class EsRagConfig{
    @Autowired
    ZhiPuAiEmbeddingModel zhiPuAiEmbeddingModel;
    @Bean
    public RestClient restClient() {
        return RestClient.builder(new HttpHost("192.168.116.6", 9200, "http")) // 配置 ES 地址和端口
                .build();
    }
    @Bean
    public ElasticsearchVectorStore elasticsearchVectorStore
            (RestClient restClient,
             ZhiPuAiEmbeddingModel zhiPuAiEmbeddingModel) {
        ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
        options.setIndexName("snail_agent"); // 设置索引名称
        options.setSimilarity(SimilarityFunction.cosine);
        options.setDimensions(2048); // 设置向量维度
        options.setEmbeddingFieldName("embedding");
        return ElasticsearchVectorStore.builder(restClient, zhiPuAiEmbeddingModel)
                .options(options)
                .initializeSchema(true) // 没有索引时自动创建索引
                .batchingStrategy(new TokenCountBatchingStrategy())//批处理策略
                .build();
    }
}